package com.sunny.k_merge;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.PriorityQueue;

//利用堆来取小值
public class KMerge2 {


    public static void main(String[] args) {

        KMerge1 kMerge1 = new KMerge1();
        ArrayList<Integer> list1 = new ArrayList<>(Arrays.asList(1,3,5,7,8,9));
        ArrayList<Integer> list2 = new ArrayList<>(Arrays.asList(1,2,3,5,6,10,11));
        ArrayList<Integer> list3 = new ArrayList<>(Arrays.asList(1,1,3,3,8,8,9,10));
        ArrayList<ArrayList<Integer>> lists = new ArrayList<>();
        lists.add(list1);
        lists.add(list2);
        lists.add(list3);
        ArrayList<Integer> res = kMerge1.sort(lists);
        res.forEach(System.out::println);

    }


    public ArrayList<Integer> sort(ArrayList<ArrayList<Integer>> arr){

        ArrayList<Integer> res = new ArrayList<>();

        PriorityQueue<Node> pq = new PriorityQueue<>(arr.size());

        int[] ind = new int[arr.size()];
        int n = 0;
        for(int i=0;i<arr.size();i++){
            n += arr.get(i).size();
        }
        for(int i=0;i<arr.size();i++){
            pq.offer(new Node(arr.get(i).get(ind[i]++),i));
        }
        int cnt = 0;
        while (cnt<n){
            Node tmp = pq.poll();
            res.add(tmp.val);
            ind[tmp.num]++;
            if(ind[tmp.num] == arr.get(tmp.num).size()){
                pq.offer(new Node(Integer.MAX_VALUE, -1));
            }else{
                pq.offer(new Node(arr.get(tmp.num).get(ind[tmp.num]), tmp.num));
            }
            cnt++;
        }
        return res;
    }



}

class Node implements Comparable<Node>{
    int val;
    int num;

    public Node(int val, int num) {
        this.val = val;
        this.num = num;
    }

    @Override
    public int compareTo(Node o) {
        return this.val-o.val;
    }
}